library(data.table)
library(mefa)
library(doParallel)
library(emayili)
library(magrittr)



source('helper_functions.r')


runBootstrap <- function(config) {
  
  N_BOOTSTRAPS = 100
  SAMPLE_SIZE = .1
  
  for(nn in 1:N_BOOTSTRAPS) {
    
    # Load the estimation
    print(paste0('Loading estimation result ',config))
    estimation <- loadEstimation(config) 
    
    
    # Create saving directory
    dir.create(paste0('estimation_results/',estimation$pars$filename,'_bootstrap/'))
    
    # Load data and create raw draws
    d.in <- fread('../data/final/estimation_data_sep9.csv')
    estimation.data <- estimation$market.data
    estimation.data <- estimation.data[,names(d.in),with=F]
    
    all.markets <- unique(estimation.data$MARKET_ID)
    set.seed(as.numeric(Sys.time()))
    sampled.ids = unique(sample(all.markets,size = floor(length(all.markets) * SAMPLE_SIZE),replace = F))  
    sampled.data = estimation.data[MARKET_ID %in% sampled.ids]
    runBootstrappedEstimation(estimation,sampled.data,as.integer(Sys.time()))
  }
  
}


runBootstrappedEstimation <- function(estimation.results,market.data,id) {
  
  pars <- estimation.results$pars
  
  pars$CORES <- 30
  
  print(nrow(market.data))
  markets.data <- market.data[YEAR %in% pars$EST_YEARS]
  print(nrow(markets.data))
  
  
  write.table(file = paste0('estimation_results/',pars$filename,'_bootstrap/bootstrap_data_',id,'.csv'),x = markets.data,row.names=F,col.names=T,sep=',')
  
  # Create instrumented interest rate
  rate.model <- lm(RATE ~ JUMBO * (D60.L1 ) + as.factor(paste0(YEAR,CBSA))   , data = markets.data,weights=markets.data$MARKET_SIZE)
  markets.data[,RATE_INST := predict(rate.model,markets.data)]
  
  # Set initial guess and raw draws. Initial guess is just the solution of the model
  new_mx <- param_vec_to_matrix(x = estimation.results$solution)
  pars$A <- new_mx$A
  pars$B <- new_mx$B
  pars$S <- new_mx$S
  raw.draws <- createRawDraws(pars)
  delta.init <<- NULL
  
  # Pre-loop
  # Calculate weights (if needed)
  # Initialize delta and drop some unsolveable markets.
  
  temp.pars <- pars
  temp.pars$BLP_TOL = 1e-10
  temp.pars$BLP_ATTEMPTS = 1000
  contraction <- solveAllMarkets(markets.data,raw.draws,temp.pars,delta.init,cores = temp.pars$CORES)
  contraction <- contraction[converged == 1]
  markets.data <- markets.data[MARKET_ID %in% contraction$MARKET_ID]
  delta.init <<- contraction[,c('MARKET_ID','j','delta')]
  
  moment_weights <- pars$moment_weights
  
  # Calculate GMM objective funcion
  gmm_objective <- function(x,diag = F) {
    
    
    print(x)
    
    # 1.  Deal parameters
    new_mx <- param_vec_to_matrix(x)
    pars$A <- new_mx$A
    pars$B <- new_mx$B
    pars$S <- new_mx$S
    
    # 2. Recover deltas
    contraction <- solveAllMarkets(markets.data,raw.draws,pars,delta.init,cores = pars$CORES)
    
    # 3. Update the initial guesses for the good ones.
    next.init <- contraction
    next.init[converged == 0,delta := mean(next.init[converged == 1]$delta,na.rm=T)]
    delta.init <<- next.init[,c('MARKET_ID','j','delta')]
    
    # 4. Merge the deltas to the main marketdata
    market.data.merged <- merge(markets.data,contraction,by=c('MARKET_ID','j'))
    
    # 5. Regress out the remaining parameters
    if(length(pars$EST_YEARS) > 1) {
      linear <- lm(delta ~ RATE_INST + JUMBO + as.factor(YEAR) * PURPOSE * LENDER,weights = market.data.merged$converged,data = market.data.merged)
    } else {
      linear <- lm(delta ~ RATE_INST + JUMBO + PURPOSE * LENDER,weights = market.data.merged$converged,data = market.data.merged)
    }
    
    # 6. Calculate moments, apply wieghts, and calculate objective function
    moments <- calculateMoments(market.data.merged,linear)
    moments <- merge(moments,moment_weights,by='moment')
    val     <- sum(moments$value^2 * moments$weight)
    
    # 7. Report diagnostics if we want
    if(pars$diag.show_linear) {
      print(linear)
    }
    if(pars$diag.show_moments) {
      moments[,contribution := value^2 * weight]
      print(moments)  
    }
    if(pars$diag.show_convergence) {
      print(paste0(sum(contraction$converged)/4,'/',nrow(contraction)/4,' markets converged. Mean iterations given convergence = ',mean(contraction[converged == 1]$attempts)))
    }
    
    
    # 8. Report actual objective function
    print(val)
    
    if(!diag) {
      return(val)
    } else {
      return(list(x = x, market.data.merged = market.data.merged, linear.model = linear, moments = moments, val = val))
    }  
    
    
    
  }
  
  # Run the optimization
  scale = 1e-1 * c(1,1,.1,.1,.1,.1,.1,.1,.1,.1,.1,.1,.1,.1)
  # Var used t obe dividing by 10, not 2.5
  result <- optim(par = estimation.results$solution+runif(length(scale))/2.5-.5/2.5, fn = gmm_objective,method = 'L-BFGS-B',lower = c(1,11,-1,-3,-4,0,-1,-3,-4,0,.01,.01,.01,.01),control = list(parscale = scale,lmm = 10))
  
  # Solve the model and return diagnostics to save
  solution <- gmm_objective(result$par,diag = T)
  
  nonLinPars <- data.table(names = c('mu_b','mu_f','pi_a_i','pi_b_i','pi_g_i','pi_f_i','pi_a_p','pi_b_p','pi_g_p','pi_f_p','sig_a','sig_b','sig_g','sig_f'),values = solution$x)
  linPars <- data.table(names = names(solution$linear.model$coefficients), values = as.numeric(solution$linear.model$coefficients))
  results <- rbind(nonLinPars,linPars)
  
  write.table(file = paste0('estimation_results/',pars$filename,'_bootstrap/bootstrap_',id,'.csv'),x = results,row.names=F,col.names=T,sep=',')
  
  
  
}



## Execute it.


# Load config parameters
config = commandArgs(trailingOnly = T)


## Run the estimation
runBootstrap(config)




